package com.baijiaoxi.common.model.concurrent;

import com.baijiaoxi.common.model.local.ThreadLocalUtil;
import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;

@Slf4j
public class BaiExecutors {
    private static final Long defaultKeepAliveTime = 60L;
    private static final TimeUnit defaultUnit = TimeUnit.SECONDS;
    public static Map<String, ThreadPoolExecutor> executorMap = new HashMap<>();

    /**
     * 有界
     * 必须指定阻塞超限处理
     * 安全
     *
     * @param poolName
     * @param corePoolSize
     * @param maximumPoolSize
     * @param maximumBlockSize 阻塞队列最大阻塞数量
     * @param handler
     * @return
     */

    public static BaiThreadPoolExecutor create(String poolName, int corePoolSize, int maximumPoolSize, int maximumBlockSize, RejectedExecutionHandler handler) {
        BaiThreadPoolExecutor baiThreadPoolExecutor = new BaiThreadPoolExecutor(poolName, corePoolSize, maximumPoolSize,
                defaultKeepAliveTime, defaultUnit, new ArrayBlockingQueue<>(maximumBlockSize),
                new BaiThreadFactory(poolName), handler);
        executorMap.put(poolName, baiThreadPoolExecutor);
        return baiThreadPoolExecutor;
    }

    /**
     * 有界
     *
     * @param poolName
     * @param corePoolSize
     * @param maximumPoolSize
     * @param maximumBlockSize
     * @return
     */

    public static BaiThreadPoolExecutor create(String poolName, int corePoolSize, int maximumPoolSize, int maximumBlockSize) {
        BaiThreadPoolExecutor baiThreadPoolExecutor = new BaiThreadPoolExecutor(poolName, corePoolSize, maximumPoolSize,
                defaultKeepAliveTime, defaultUnit, new ArrayBlockingQueue<>(maximumBlockSize), new BaiThreadFactory(poolName),
                new RejectedExecutionHandler() {
                    @Override
                    public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
                        throw new RuntimeException("系统繁忙");
                    }
                });
        executorMap.put(poolName, baiThreadPoolExecutor);
        return baiThreadPoolExecutor;
    }

    /**
     * 边界阻塞（Integer.MAX_VALUE）
     * 谨慎使用，只有在确定的短任务可以使用
     * 不安全，可能会造成oom
     *
     * @param poolName
     * @param corePoolSize
     * @param maximumPoolSize
     * @return
     */
    public static BaiThreadPoolExecutor create(String poolName, int corePoolSize, int maximumPoolSize) {
        BaiThreadPoolExecutor baiThreadPoolExecutor = new BaiThreadPoolExecutor(poolName, corePoolSize, maximumPoolSize,
                defaultKeepAliveTime, defaultUnit, new LinkedBlockingQueue<>(), new BaiThreadFactory(poolName), new RejectedExecutionHandler() {
            @Override
            public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
                throw new RuntimeException("系统繁忙");
            }
        });
        executorMap.put(poolName, baiThreadPoolExecutor);

        return baiThreadPoolExecutor;
    }

    public static BaiThreadPoolExecutor create(String poolName, int corePoolSize) {
        BaiThreadPoolExecutor baiThreadPoolExecutor = new BaiThreadPoolExecutor(poolName, corePoolSize, corePoolSize * 2, defaultKeepAliveTime, defaultUnit,
                new LinkedBlockingQueue<>(), new BaiThreadFactory(poolName), new RejectedExecutionHandler() {
            @Override
            public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
                throw new RuntimeException("系统繁忙");
            }
        });
        executorMap.put(poolName, baiThreadPoolExecutor);

        return baiThreadPoolExecutor;
    }

    public static ScheduledThreadPoolExecutor createScheduled(String poolName, int corePoolSize) {
        return createScheduled(poolName, corePoolSize, new RejectedExecutionHandler() {
            @Override
            public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
                throw new RuntimeException("系统繁忙");
            }
        });
    }

    public static ScheduledThreadPoolExecutor createScheduled(String poolName, int corePoolSize, RejectedExecutionHandler handler) {
        ScheduledThreadPoolExecutor scheduledThreadPoolExecutor =
                new ScheduledThreadPoolExecutor(corePoolSize, new BaiThreadFactory(poolName), handler);
        executorMap.put(poolName, scheduledThreadPoolExecutor);
        return scheduledThreadPoolExecutor;
    }

    /**
     * 所有线程池活跃线程数合计
     * 执行任务中的，任务执行完未回收的不算
     *
     * @return
     */
    public static Long getSumActiveCount() {
        AtomicLong count = new AtomicLong();
        executorMap.entrySet().stream().forEach(executor -> {
            count.addAndGet(executor.getValue().getActiveCount());
        });
        return count.get();
    }

    /**
     * 获取所有线程池线程数合计
     *
     * @return
     */
    public static Long getSumPoolSize() {
        AtomicLong count = new AtomicLong();
        executorMap.entrySet().stream().forEach(executor -> {
            count.addAndGet(executor.getValue().getPoolSize());
        });
        return count.get();
    }

    /**
     * 获取所有线程池阻塞数合计
     *
     * @return
     */
    public static Long getSumQueueSize() {
        AtomicLong count = new AtomicLong();
        executorMap.entrySet().stream().forEach(executor -> {
            count.addAndGet(executor.getValue().getQueue().size());
        });
        return count.get();
    }

    /**
     * 获取线程池对象
     *
     * @param name
     * @return
     */
    public static ThreadPoolExecutor getThreadPoolExecutor(String name) {
        return executorMap.get(name);
    }


    /**
     * 执行一个简单的线程
     *
     * @param runnable
     */
    public static void execute(Runnable runnable) {
        Object va = ThreadLocalUtil.get();
        Map<String, String> contextMap = MDC.getCopyOfContextMap();
        new Thread(new Runnable() {
            @Override
            public void run() {
                if (va != null)
                    ThreadLocalUtil.set(va);
                if (contextMap != null)
                    MDC.setContextMap(contextMap);
                try {
                    runnable.run();
                } catch (Exception e) {
                    log.error("线程执行错误", e);
                } finally {
                    ThreadLocalUtil.clear();
                    MDC.clear();
                }
            }
        }).start();
    }
}
